#!/usr/bin/env python3
"""
add_idx_to_jsonl.py

将 idx 字段添加到每条 JSONL 记录中（手动在脚本中设置路径与起始编号）。

在脚本顶部修改下面变量：
- INPUT_JSONL: 输入 JSONL 文件路径
- OUTPUT_JSONL: 输出 JSONL 文件路径
- START_IDX: 起始编号（例如 2000）
- OVERWRITE_EXISTING_IDX: 是否覆盖已有的 idx 字段（True/False）
- ON_JSON_DECODE_ERROR: "skip" 或 "keep"（当某行无法解析为 JSON 时：skip 表示跳过该行；keep 表示原样写回到输出文件）

运行：
python add_idx_to_jsonl.py
"""

import json
import os
import sys

# =========== 这里手动设置 ===========
INPUT_JSONL = "eval/data/mawps/1AIME2025.jsonl"
OUTPUT_JSONL = "eval/data/mawps/AIME2025.jsonl"
START_IDX = 2812
OVERWRITE_EXISTING_IDX = True   # 若为 False，则遇到已有 idx 字段时保持原样（不覆盖）
ON_JSON_DECODE_ERROR = "skip"   # 可选 "skip" 或 "keep"
# ====================================


def add_idx_stream(input_path, output_path, start_idx=0, overwrite=True, on_error="skip"):
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"输入文件不存在: {input_path}")

    if on_error not in ("skip", "keep"):
        raise ValueError("on_error 参数必须是 'skip' 或 'keep'")

    idx = int(start_idx)
    total = 0
    written = 0
    skipped = 0
    errored = 0

    with open(input_path, "r", encoding="utf-8") as fin, \
         open(output_path, "w", encoding="utf-8") as fout:
        for line_no, raw in enumerate(fin, start=1):
            raw_stripped = raw.rstrip("\n")
            if raw_stripped.strip() == "":
                # 跳过空行（不计入编号）
                continue
            total += 1
            try:
                obj = json.loads(raw_stripped)
            except Exception as e:
                errored += 1
                msg = f"[WARN] JSON decode error at line {line_no}: {e}"
                print(msg, file=sys.stderr)
                if on_error == "keep":
                    # 将原始行写回输出（未经修改）
                    fout.write(raw_stripped + "\n")
                    written += 1
                else:
                    skipped += 1
                continue

            # 确保为 dict
            if not isinstance(obj, dict):
                # 若不是对象（例如列表/数字），跳过或保留原样
                print(f"[WARN] line {line_no} parsed to non-dict ({type(obj)}). Skipping.", file=sys.stderr)
                skipped += 1
                continue

            # 决定是否写入 idx
            if not overwrite and ('idx' in obj):
                # 不覆盖已有 idx，直接写回原对象
                fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
                written += 1
            else:
                obj['idx'] = idx
                idx += 1
                fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
                written += 1

    print("===== 完成 =====")
    print(f"输入文件: {input_path}")
    print(f"输出文件: {output_path}")
    print(f"总遍历行数（含空/不可解析行）: {total}")
    print(f"成功写入记录数: {written}")
    print(f"跳过记录数（不可解析或空行）: {skipped}")
    print(f"JSON 解码错误行数: {errored}")
    print(f"下一个 idx 将从: {idx} 开始（若继续处理其它文件可接着用）")


if __name__ == "__main__":
    add_idx_stream(
        INPUT_JSONL,
        OUTPUT_JSONL,
        start_idx=START_IDX,
        overwrite=OVERWRITE_EXISTING_IDX,
        on_error=ON_JSON_DECODE_ERROR
    )
